#!/usr/bin/env python
# -*- coding:utf-8 -*-

""" 
@author: zhaoqifa@baidu.com
@license: Copyright (c) 2013 Baidu.com, Inc. All Rights Reserved
"""
import threading
from threading import Thread as Thread
#from lib.kthread import KThread

class Task:
    def __init__(self,callable,*args,**kwargs):
        self.callable = callable
        self.args = args
        self.kwargs = kwargs

class PCTaskPool:
    """
    Task pool for producer and consumer model
    """
    def __init__(self,producer_num,consumer_num,consumer_func):
        self.tasks = []
        self.products = []

        self.producer_num = producer_num
        self.producers = []
        
        self.consumer_num = consumer_num
        self.consumer_func = consumer_func
        self.consumers = []

        self.producer_lock = threading.Lock()
        self.consumer_lock = threading.Lock()

    def create_workers(self):
        for i in xrange(self.producer_num):
            self.producers.append(Thread(target=self.product))
        for i in xrange(self.consumer_num):
            self.consumers.append(Thread(target=self.consume))
    
    def add_task(self,task):
        """
        All tasks should be submitted before taskpool start running, no need to use lock
        """
        #self.producer_lock.acquire()
        self.tasks.append(task)
        #self.producer_lock.release()

    def fetch_task(self):
        task = None
        self.producer_lock.acquire()
        if len(self.tasks) != 0:
            task = self.tasks.pop(0)
        self.producer_lock.release()
        return task

    def add_product(self,product):
        self.consumer_lock.acquire()
        self.products.append(product)
        self.consumer_lock.release()

    def fetch_product(self):
        product = None
        self.consumer_lock.acquire()
        if len(self.products) != 0:
            product = self.products.pop(0)
        self.consumer_lock.release()
        return product
        
    def product(self):
        task = self.fetch_task()
        while task is not None:
            p = task.callable(*task.args,**task.kwargs)
            self.add_product(p)
            task = self.fetch_task()
    
    def has_product_in_future(self):
        if len(self.tasks) > 0:
            return True
        for producer in self.producers:
            if producer.isAlive():
                return True
        return False
    
    def consume(self):
        product = self.fetch_product()
        while product is not None or self.has_product_in_future():
            if product is not None:
                self.consumer_func(product)
            product = self.fetch_product()
        
    def start(self):
        if len(self.tasks) == 0:
            return 0
        if self.producer_num > len(self.tasks):
            self.producer_num = len(self.tasks)
        self.create_workers()
    
        for producer in self.producers:
            producer.start()
        for consumer in self.consumers:
            consumer.start()
        for producer in self.producers:
            producer.join()
        for consumer in self.consumers:
            consumer.join()
        
